import os
import requests
from typing import Any, Dict, List, Optional, Tuple
from pydantic import ValidationError
import logging

# Import the models that ServerHelper needs
from .model.task import TaskSpec, TaskStatus

logger = logging.getLogger(__name__)


class ServerHelper:
    def __init__(
        self,
        base_url: str = "http://localhost:9010",
    ):
        self.base_url = base_url.rstrip("/")

    def _parse_and_validate(self, data, model_cls):
        """Helper method to parse JSON data and validate against a model class.

        Args:
            data: The data to parse and validate
            model_cls: The Pydantic model class to validate against

        Returns:
            Parsed model instance(s) or raw JSON data
        """
        if model_cls is not None and data is not None:
            try:
                if isinstance(data, list):
                    return [model_cls.model_validate(item) for item in data]
                return model_cls.model_validate(data)
            except ValidationError:
                # If model validation fails, return raw data
                pass
        return data

    def _get(self, path: str, params: Optional[Dict[str, Any]] = None, model_cls=None):
        """Helper method for GET requests.

        Returns:
            Tuple containing:
            - Parsed model instance(s) or raw JSON data
            - HTTP status code
        """
        response = requests.get(f"{self.base_url}/{path.lstrip('/')}", params=params)
        status_code = response.status_code

        # Raise for status if not 2xx
        if not 200 <= status_code < 300:
            response.raise_for_status()

        # Try to parse JSON, return raw response if parsing fails
        try:
            data = response.json()
        except ValueError:
            return response.text, status_code

        parsed_data = self._parse_and_validate(data, model_cls)
        return parsed_data, status_code

    def _post(self, path: str, data: Any, model_cls=None):
        """Helper method for POST requests.

        Returns:
            Tuple containing:
            - Parsed model instance or raw JSON data
            - HTTP status code
        """
        response = requests.post(f"{self.base_url}/{path.lstrip('/')}", json=data)
        status_code = response.status_code

        # Raise for status if not 2xx
        if not 200 <= status_code < 300:
            response.raise_for_status()

        # Try to parse JSON, return raw response if parsing fails
        try:
            response_data = response.json()
        except ValueError:
            return response.text, status_code

        parsed_data = self._parse_and_validate(response_data, model_cls)
        return parsed_data, status_code

    def _put(
        self,
        path: str,
        data: Any,
        model_cls=None,
        params: Optional[Dict[str, Any]] = None,
    ):
        """Helper method for PUT requests.

        Returns:
            Tuple containing:
            - Parsed model instance or raw JSON data
            - HTTP status code
        """
        response = requests.put(
            f"{self.base_url}/{path.lstrip('/')}", json=data, params=params
        )
        status_code = response.status_code

        # Raise for status if not 2xx
        if not 200 <= status_code < 300:
            response.raise_for_status()

        # Try to parse JSON, return raw response if parsing fails
        try:
            response_data = response.json()
        except ValueError:
            return response.text, status_code

        parsed_data = self._parse_and_validate(response_data, model_cls)
        return parsed_data, status_code

    def create_task_spec(self, task_spec: TaskSpec) -> Tuple[TaskSpec, bool]:
        """Create a new TaskSpec in the server.

        Returns a tuple of (TaskSpec, bool) where the bool is True if the task spec was created,
        and False if it already exists.
        """
        result, status_code = self._post(
            "task_specs/", task_spec.model_dump(), TaskSpec
        )
        # Determine if the task was created or already existed based on status code
        is_created = status_code == 200
        return result, is_created

    def get_task_spec(self, task_hash: str) -> TaskSpec:
        """Get a TaskSpec by its task_hash."""
        result, _ = self._get(f"task_specs/{task_hash}", model_cls=TaskSpec)
        return result

    def list_task_specs(self, offset: int = 0, limit: int = 10) -> list[TaskSpec]:
        """Get a list of TaskSpecs with pagination."""
        result, _ = self._get(
            "task_specs/", params={"offset": offset, "limit": limit}, model_cls=TaskSpec
        )
        return result

    def list_named_task_specs(self, offset: int = 0, limit: int = 10) -> list[TaskSpec]:
        """Get a list of TaskSpecs that have non-empty names with pagination.
        
        Args:
            offset: Number of records to skip (default: 0)
            limit: Maximum number of records to return (default: 10, max: 100)
            
        Returns:
            List of TaskSpec objects that have non-empty names
            
        Raises:
            requests.HTTPError: If the API call fails
        """
        result, _ = self._get(
            "named_task_specs/", params={"offset": offset, "limit": limit}, model_cls=TaskSpec
        )
        return result

    def set_task_status(self, task_hash: str, status: TaskStatus) -> TaskStatus:
        """Update the status of a task by its task_hash.

        Args:
            task_hash: The hash of the task to update
            status: The new status to set

        Returns:
            The updated TaskStatus

        Raises:
            ValueError: If status is not a valid TaskStatus value
        """
        # Validate that status is a valid TaskStatus
        if not isinstance(status, TaskStatus):
            raise ValueError(f"Status must be a TaskStatus enum, got {type(status)}")
        data, _ = self._put(
            f"task_status/{task_hash}", None, params={"status": status.value}
        )
        return status

    def get_ready_tasks(self, limit: int = 1) -> List[TaskSpec]:
        """
        Find and return a list of tasks that are ready to run.

        A task is ready if:
        1. It is in PENDING status
        2. All its dependencies (in file_params and output_params) are in COMPLETED status

        If any dependency is in FAILED status, the task will be marked as FAILED.

        Note: The API may return either TaskSpec objects or task_hash strings.
        This method handles both cases and always returns TaskSpec objects.

        Args:
            limit: Maximum number of ready tasks to return (1-10)

        Returns:
            List of TaskSpec objects that are ready to run
        """
        if limit < 1 or limit > 10:
            raise ValueError("Limit must be between 1 and 10")

        result, status_code = self._get(
            "task_ready", params={"limit": limit}, model_cls=TaskSpec
        )

        # If no tasks are ready, return empty list
        if status_code == 204 or not result:
            return []

        # The API might return a list of task hashes or TaskSpec objects
        if isinstance(result, list) and result and isinstance(result[0], str):
            # Convert task hashes to TaskSpec objects
            task_specs = []
            for task_hash in result:
                try:
                    task_spec = self.get_task_spec(task_hash)
                    task_specs.append(task_spec)
                except requests.HTTPError:
                    # Skip tasks that couldn't be retrieved
                    continue
            return task_specs

        return result

    def acquire_task(self, task_hash: str) -> Optional[TaskSpec]:
        """
        Find and return a task that is ready to run, and mark it as RUNNING.

        Returns None if no task is ready or could not be acquired.
        """
        try:
            acquired_data, acquired_status = self._post(
                f"acquire_task/{task_hash}", {}, TaskSpec
            )

            # If task was not successfully acquired, return None
            if acquired_status == 204 or acquired_data is None:
                return None

            # Get the full task spec to ensure we have the most up-to-date information
            return self.get_task_spec(task_hash)
        except requests.HTTPError:
            # If there was an error acquiring the task, return None
            return None

    def _upload_file(
        self, path: str, object_path: str, content_type: str | None = None
    ):
        """Upload a file to the server."""
        # check file exists
        if not os.path.exists(path):
            raise FileNotFoundError(f"File {path} does not exist")
        # check file is not a directory
        if os.path.isdir(path):
            raise IsADirectoryError(f"File {path} is a directory")
        with open(path, "rb") as f:
            start_bytes = 0
            while True:
                chunk = f.read(2**16)
                response = requests.put(
                    f"{self.base_url}/file/{object_path}",
                    data=chunk,
                    params={"start_bytes": start_bytes, "end_bytes": start_bytes + len(chunk)},
                    headers={"Content-Type": content_type} if content_type else None,
                )
                if response.status_code != 200:
                    raise requests.HTTPError(
                        f"Failed to upload file {path} to {object_path}: {response.text}"
                    )
                start_bytes += len(chunk)
                if len(chunk) == 0:
                    break
                

    def upload_file(self, path: str, object_path: str):
        """Upload a file to the server."""
        logger.info(f"Uploading file {path} to {object_path}")
        self._upload_file(path, object_path)

    def upload_task_file(self, path: str, task_hash: str, object_name: str):
        """Upload a file to the server."""
        # remove the first strip_prefix parts of the path
        object_path = os.path.join(task_hash, "file", object_name)
        self._upload_file(path, object_path)

    def upload_task_output(self, path: str, task_hash: str, object_name: str):
        """Upload a file to the server."""
        # remove the first strip_prefix parts of the path
        object_path = os.path.join(task_hash, "output", object_name)
        self._upload_file(path, object_path, content_type="application/pickle")

    def upload_task_log(self, path: str, task_hash: str, object_name: str):
        """Upload a file to the server."""
        # remove the first strip_prefix parts of the path
        object_path = os.path.join(task_hash, "log", object_name)
        self._upload_file(path, object_path, content_type="application/json")

    def _get_chunk(self, object_name: str, offset: int, chunk_size: int) -> bytes:
        response = requests.get(
            f"{self.base_url}/file/{offset}/{offset + chunk_size}/{object_name}"
        )
        return response.content

    def download_file(self, object_name, target_path):
        offset = 0
        chunk_size = 2**20
        os.makedirs(os.path.dirname(target_path), exist_ok=True)
        with open(target_path, "wb") as f:
            while True:
                chunk = self._get_chunk(object_name, offset, chunk_size)
                if chunk is None or len(chunk) == 0:
                    break
                f.write(chunk)
                offset += len(chunk)
        return offset

    def list_files(self, task_hash: str|None = None, _type: str|None = None) -> list[str]:
        """List all files in the task directory."""
        params = {}
        if task_hash is not None:
            params["task_hash"] = task_hash
        if _type is not None:
            params["_type"] = _type
        result, _ = self._get(f"file_list", params=params)
        return result

    def delete_task(self, task_hash: str) -> List[str]:
        """
        Delete a task and all tasks that depend on it.

        This deletes task records and dependency records from the database.
        Additionally, removes any files stored in Minio for all deleted tasks.

        Args:
            task_hash: The hash of the task to delete

        Returns:
            List of deleted task hashes

        Raises:
            requests.HTTPError: If the API call fails
        """
        try:
            # Delete task records from the API server
            result, _ = self._delete(f"task/{task_hash}")

            # Get list of deleted tasks
            deleted_tasks = result.get("deleted_tasks", [])

            # Delete all files associated with the deleted tasks
            for task in deleted_tasks:
                # Delete all files for this task
                self._delete_task_files(task)

            return deleted_tasks
        except Exception as e:
            logger.exception(e)
            raise

    def reset_task(self, task_hash: str) -> List[str]:
        """
        Reset a task and all tasks that depend on it.

        This sets the status of the tasks back to PENDING.
        Additionally, removes any files stored in Minio for all reset tasks.
        The task records and dependency records are not removed.

        Args:
            task_hash: The hash of the task to reset

        Returns:
            List of reset task hashes

        Raises:
            requests.HTTPError: If the API call fails
        """
        try:
            # Reset task status via the API server
            reset_tasks, _ = self._post(f"task/{task_hash}/reset", {})
            # Delete all files associated with the reset tasks
            for task in reset_tasks:
                # Delete all files for this task
                self._delete_task_files(task)

            return reset_tasks
        except Exception as e:
            logger.exception(e)
            raise

    def _delete_task_files(self, task_hash: str):
        """
        Delete all files associated with a task from Minio.

        Args:
            task_hash: The hash of the task
        """
        try:
            # Delete all files with the task_hash prefix, 
            self._delete(f"file/{task_hash}/")
        except Exception as e:
            logger.warning(f"Failed to delete files for task {task_hash}: {str(e)}")

    def get_dependencies(self, task_hash: str) -> List[str]:
        """
        Get direct dependencies of a task.

        Args:
            task_hash: The hash of the task to get dependencies for

        Returns:
            List of task hashes that this task directly depends on
        """
        try:
            result, _ = self._get(f"task/{task_hash}/dependencies")
            return result
        except Exception as e:
            logger.exception(e)
            raise

    def _delete(self, path: str, params: Optional[Dict[str, Any]] = None):
        """Helper method for DELETE requests.

        Returns:
            Tuple containing:
            - Parsed model instance(s) or raw JSON data
            - HTTP status code
        """
        response = requests.delete(f"{self.base_url}/{path.lstrip('/')}", params=params)
        status_code = response.status_code

        # Raise for status if not 2xx
        if not 200 <= status_code < 300:
            response.raise_for_status()

        # Try to parse JSON, return raw response if parsing fails
        try:
            data = response.json()
        except ValueError:
            return response.text, status_code

        return data, status_code 